import numpy as np
from common.python.common import *

class ImageEncoder:
    def __init__(self):
        pass

    def encode_image_and_save(self, stream, output_path, image, img_info):
        attr = sdk.lyn_jenc_attr_t()
        attr.input_format = img_info.output.output_fmt
        attr.width = img_info.output.width
        attr.height = img_info.output.height
        attr.jpeg_quality = 95

        info = sdk.lyn_jenc_info_t()
        info.params = attr
        info.input = image
        info.output.size = image.size
        info.output.data, ret = sdk.lyn_malloc(info.output.size)
        error_check(ret, "lyn_malloc")

        error_check(sdk.lyn_jpeg_encode_async(stream, info))
        error_check(sdk.lyn_synchronize_stream(stream))

        packet = sdk.lyn_packet_t()
        packet.data = info.output.data
        packet.size = info.output.size
        outdata_size, ret = sdk.lyn_enc_get_remote_packet_valid_size(packet)
        error_check(ret, "lyn_enc_get_remote_packet_valid_size")

        outdata = np.zeros(outdata_size, dtype=np.uint8)
        outdata_ptr = sdk.lyn_numpy_to_ptr(outdata)
        error_check(
            sdk.lyn_memcpy(
                outdata_ptr,
                packet.data,
                outdata_size,
                sdk.lyn_memcpy_dir_t.ServerToClient,
            )
        )

        with open(output_path, "wb") as f:
            f.write(outdata.tobytes())

        error_check(sdk.lyn_free(info.output.data))
